# -*- coding: utf-8 -*-

import tensorflow as tf
import numpy as np
import data1.config, data2.config, data3.config, data4.config
import data1.preprocessing, data2.preprocessing, data3.preprocessing, data4.preprocessing
import utils
import threading


class threadsafe_generator:
    """Takes an generator and makes it thread-safe by
    serializing call to the `next` method of given generator.
    """

    def __init__(self, gen):
        self.gen = gen
        self.lock = threading.Lock()

    def __iter__(self):
        return self

    def __next__(self):
        with self.lock:
            return next(self.gen)

    def send(self, data):
        with self.lock:
            return self.gen.send(data)


def generator_distinguish_one_char(model_name, one_hot2char_func, char_image_width, char_image_height, label_size):
    """
    创建一个识别单个字符的生成器，该生成器可以不断识别单个字符
    <code>
        gen = generator_distinguish_one_char(model_name, one_hot2char_func, char_image_width, char_image_height, label_size)
        # start generator
        next(gen)
        result = gen.send(char_image)
        print(result)
    </code>
    :param model_name:
    :param one_hot2char_func:
    :param char_image_width:
    :param char_image_height:
    :param label_size:
    :return: None
    """
    saver = tf.train.import_meta_graph(model_name + '.meta')
    with tf.Session() as sess:
        saver.restore(sess, model_name)
        char = None
        while True:
            char_image = yield char
            out_one_hot = sess.run('output_label:0', feed_dict={
                'input_images:0': char_image.reshape(1, char_image_width, char_image_height),
                'input_labels:0': np.zeros((1, label_size)),
                'keep_prob:0': 1.0
            })
            char = one_hot2char_func(out_one_hot)


class Distinguish:
    """
    验证码识别类
    :param model_name tensorflow 模型的路径 + 名称
    :param one_hot2char_func 将 one_hot 数据转换为字符的函数
    :param char_image_width 单个字符图片的宽度
    :param char_image_height 单个字符图片的高度
    :param label_size 标签的大小(分类的个数)
    """

    def __init__(self, model_name, one_hot2char_func, char_image_width, char_image_height, label_size,
                 image_preprocessor):

        self.distinguish_one_char_generator = threadsafe_generator(generator_distinguish_one_char(model_name=model_name,
                                                                                                  one_hot2char_func=one_hot2char_func,
                                                                                                  char_image_width=char_image_width,
                                                                                                  char_image_height=char_image_height,
                                                                                                  label_size=label_size))

        # start generator
        next(self.distinguish_one_char_generator)
        self.image_preprocessor = image_preprocessor

    def distinguish(self, img, is_gray=False):
        """
        识别图片
        :param img: 带识别的图片
        :param is_gray: 是否为灰度图像
        :return: 识别后的字符串
        """
        if not is_gray:
            img = utils.load_image_by_bytes(img)
        char_image_array = self.image_preprocessor(img)
        yzm_str = ''
        for char_image in char_image_array:
            yzm_str += self.distinguish_one_char_generator.send(char_image)
        return yzm_str

    def distinguish_by_filename(self, filename):
        return self.distinguish(utils.load_image_by_filename(filename))


class DistinguishFactory:
    def __init__(self):
        self.objs = {}
        self.lock = threading.Lock()

    def get(self, category):
        with self.lock:
            if not category in self.objs:
                if category == 'data1':
                    model_name = data1.config.model_name
                    one_hot2char_func = data1.preprocessing.one_hot2char
                    char_image_width = data1.config.char_image_width
                    char_image_height = data1.config.char_image_height
                    label_size = data1.config.label_size
                    preprocessing_func = data1.preprocessing.preprocessing
                elif category == 'data2':
                    model_name = data2.config.model_name
                    one_hot2char_func = data2.preprocessing.one_hot2char
                    char_image_width = data2.config.char_image_width
                    char_image_height = data2.config.char_image_height
                    label_size = data2.config.label_size
                    preprocessing_func = data2.preprocessing.preprocessing
                elif category == 'data3':
                    model_name = data3.config.model_name
                    one_hot2char_func = data3.preprocessing.one_hot2char
                    char_image_width = data3.config.char_image_width
                    char_image_height = data3.config.char_image_height
                    label_size = data3.config.label_size
                    preprocessing_func = data3.preprocessing.preprocessing
                elif category == 'data4':
                    model_name = data4.config.model_name
                    one_hot2char_func = data4.preprocessing.one_hot2char
                    char_image_width = data4.config.char_image_width
                    char_image_height = data4.config.char_image_height
                    label_size = data4.config.label_size
                    preprocessing_func = data4.preprocessing.preprocessingWithPadding
                else:
                    model_name = data1.config.model_name
                    one_hot2char_func = data1.preprocessing.one_hot2char
                    char_image_width = data1.config.char_image_width
                    char_image_height = data1.config.char_image_height
                    label_size = data1.config.label_size
                    preprocessing_func = data1.preprocessing.preprocessing

                self.objs[category] = Distinguish(model_name, one_hot2char_func, char_image_width, char_image_height,
                                                  label_size,
                                                  preprocessing_func)
            obj = self.objs[category]
        return obj


distinguish_factory = DistinguishFactory()
